import json
import time
from typing import List, Dict

from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_ollama import OllamaLLM
from tqdm import tqdm

from model import Discipline, DisciplineDetail

OLLAMA_URL = "http://localhost:11434"
MODEL_NAME = "qwen2.5:32b"
PROMPT_TEMPLATE = """
你是一个专门提供学术学科综合信息的AI助手。当给定一个学科名称及其父学科名称时，你的任务是生成一个详细的JSON响应。请注意：

1. 如果对某些信息不确定，应明确标注"待考证"而不是编造
2. 对于历史信息，应尽可能提供可验证的时间节点和重要事件
3. 对于代表著作，仅列举公认的经典著作，避免争议性或次要作品
4. 如果信息不足，宁可从略，也不要填充可疑内容

响应信息包含：

1. 学科的简要介绍（基于公认的学科定义和范畴）
2. 学科的发展历史（重点标注可考证的重要时间节点）
3. 这个学科的代表著作（仅包含公认的经典著作）
   - 著作名称（包含原名和译名，如适用）
   - 作者（包含生卒年份，如可考）
   - 著作的核心贡献（避免主观评价）

请按照以下结构格式化你的响应为有效的JSON对象：

{format_instructions}

响应要求：
1. 所有内容必须基于可靠的学术来源
2. 对于存疑信息，使用"据考证"、"一般认为"等措辞
3. 如果某个信息点缺乏可靠来源，应在JSON中标注为null或"信息待考证"
4. 代表著作应至少包含一部公认的经典著作，如无法确定则说明"暂缺可靠记载"
5. 格式正确，是有效的JSON
6. 所有内容均使用中文
7. 避免使用模糊或笼统的时间表述，如"很早以前"、"近些年来"等

输入: 
[学科名称]: {discipline_name}
[父学科名称]: {parent_discipline_name}

在生成响应时，如果发现信息不足或可信度存疑，请在相应字段中标注说明，而不是试图补充可能不准确的信息。
"""


def generate_detail(disciplines: List[Discipline], retry_times: int = 3) -> Dict[str, DisciplineDetail]:
    parser = PydanticOutputParser(pydantic_object=DisciplineDetail)
    template = PromptTemplate(
        template=PROMPT_TEMPLATE,
        input_variables=["discipline_name", "parent_discipline_name"],
        partial_variables={
            "format_instructions": parser.get_format_instructions()
        }
    )

    llm = OllamaLLM(
        base_url=OLLAMA_URL,
        model=MODEL_NAME,
        temperature=0.1,
        num_predict=4096,
        num_ctx=1024,
    )
    failed_data = []
    result = {}
    # 初始化进度条
    total_items = len([d for d in disciplines if d.super_disciplines != ""])
    pbar = tqdm(total=total_items, desc="Generating details")

    start_time = time.time()
    processed_items = 0

    for d in disciplines:
        if not d.super_disciplines:
            continue

        if d.name not in ['瓜果学', '煤气能', '化工机械与设备', '机场工程', '西班牙语、葡萄牙语', '梵语、印地语、乌尔都语、僧伽罗语', '印度尼西亚语、菲律宾国语、马来语', '外国经济思想史', '合同战役学']:
            continue

        detail: DisciplineDetail | None = None
        for i in range(retry_times):
            try:
                prompt = template.format(
                    discipline_name=d.name,
                    parent_discipline_name=d.super_disciplines[-1]
                )
                detail = parser.parse(llm.invoke(prompt))
                break  # 如果成功，跳出重试循环
            except Exception as e:
                print(e)
                if i == retry_times - 1:  # 如果是最后一次重试
                    failed_data.append(d.name)

        if detail is not None:
            result[d.name] = detail

        processed_items += 1
        pbar.update(1)

        # 估算剩余时间
        elapsed_time = time.time() - start_time
        items_left = total_items - processed_items
        if processed_items > 0:
            avg_time_per_item = elapsed_time / processed_items
            estimated_time_left = avg_time_per_item * items_left
            pbar.set_postfix({"Est. Time Left": f"{estimated_time_left:.2f}s"})

    pbar.close()
    print("Failed items:", failed_data)
    return result


if __name__ == '__main__':
    with open("../assets/base_data.json", 'r') as f:
        base_data = json.load(f)
    base_data = [Discipline.model_validate_json(d) for d in base_data]
    detail_data = generate_detail(base_data)
    detail_data = {k: v.model_dump_json() for k, v in detail_data.items()}
    with open("../assets/detail_data.json", "w", encoding="utf-8") as f:
        json.dump(detail_data, f)
